#include "pch.h"
#include "../Base.h"

#define COMMID  0x12345678;


typedef struct _IO_STATUS_BLOCK {
	union {
		ULONG Status;
		PVOID    Pointer;
	};
	ULONG_PTR Information;
} IO_STATUS_BLOCK, * PIO_STATUS_BLOCK;

typedef   ULONG(*pfnNtQueryInformationFile)(
	HANDLE                 FileHandle,
	PIO_STATUS_BLOCK       IoStatusBlock,
	PVOID                  FileInformation,
	ULONG                  Length,
	ULONG FileInformationClass
	);


typedef ULONG64(_fastcall* pfnNtConvertBetweenAuxiliaryCounterAndPerformanceCounteryang)(
	char a1,
	PVOID a2,
	ULONG64* a3,
	ULONG64* a4);

pfnNtQueryInformationFile gNtQueryInformationFile = NULL;
pfnNtConvertBetweenAuxiliaryCounterAndPerformanceCounteryang  gNtConvertBetweenAuxiliaryCounterAndPerformanceCounter = NULL;
HANDLE ghFile = NULL;
BOOLEAN DriverCommInit() {
	static USHORT osBuildNumber;

	if (osBuildNumber) {
		return TRUE;
	}
	osBuildNumber = MyGetVersion();

	HMODULE hModule = GetModuleHandleA("ntdll.dll");
	if (osBuildNumber == 7600 || osBuildNumber == 7601) {
		gNtQueryInformationFile = (pfnNtQueryInformationFile)GetProcAddress(hModule, "NtQueryInformationFile");
		char bufPath[MAX_PATH] = { 0 };
		GetTempPathA(MAX_PATH, bufPath);
		strcat_s(bufPath, "\\1.txt");
		ghFile = CreateFileA(bufPath, FILE_ALL_ACCESS, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, CREATE_ALWAYS,
			FILE_ATTRIBUTE_NORMAL, FALSE);
		if (ghFile || ghFile == INVALID_HANDLE_VALUE) {
			MessageBoxA(NULL, "ͨ�ų�ʼ��ʧ��", "��ʾ", MB_OK);
		}
	}
	else {
		gNtConvertBetweenAuxiliaryCounterAndPerformanceCounter = (pfnNtConvertBetweenAuxiliaryCounterAndPerformanceCounteryang)
			GetProcAddress(hModule, "NtConvertBetweenAuxiliaryCounterAndPerformanceCounter");
	}

	return TRUE;

}

BOOLEAN DriverCommWin7(ULONG cmd, PVOID inData, SIZE_T size) {
	char buf[0xE0] = { 0 };
	PCommPackage package = (PCommPackage)buf;
	package->id = COMMID;
	package->cmd = cmd;
	package->inData = inData;
	package->inSize = size;
	package->retStatus = -1;
	IO_STATUS_BLOCK is = { 0 };

	gNtQueryInformationFile(ghFile, &is, buf, sizeof(buf), 0x34);
	if (package->retStatus == 0) {
		return TRUE;
	}
	return FALSE;
}

BOOLEAN DriverCommWin10(ULONG cmd, PVOID inData, SIZE_T size) {
	CommPackage package;
	package.id = COMMID;
	package.cmd = cmd;
	package.inData = inData;
	package.inSize = size;
	package.retStatus = -1;
	IO_STATUS_BLOCK is = { 0 };
	PCommPackage data = &package;
	ULONG64 validBuf = 0;
	gNtConvertBetweenAuxiliaryCounterAndPerformanceCounter(1, (PVOID)&data, &validBuf, NULL);
	if (package.retStatus == 0) {
		return TRUE;
	}
	return FALSE;
}

BOOLEAN DriverComm(ULONG cmd, PVOID inData, SIZE_T size) {

	if (DriverCommInit()) {
		USHORT BuildNumber = MyGetVersion();
		HMODULE hModule = GetModuleHandleA("ntdll.dll");
		if (BuildNumber == 7600 || BuildNumber == 7601) {
			return DriverCommWin7(cmd,inData,size);
		}
		else {
			
			return DriverCommWin10(cmd, inData, size);
		}
	}

	return FALSE;
}

ULONG64 DriverGetModuleInfo(HANDLE pid, char* moduleName)
{
	ModuleInfo info;
	info.pid = (ULONG64)pid;
	info.moduleName = (ULONG64)moduleName;
	if (DriverComm(CMD_GET_MOUDULE, &info, sizeof(info))) {
		if (info.moduleBase) {
			return info.moduleBase;
		}
	}
	return 0;
}

BOOLEAN DriverReadProcessMemory(HANDLE pid, PVOID address,PVOID buffer,SIZE_T size)
{
	READWRITEINFO info;
	info.pid = (ULONG64)pid;
	info.buffer = (ULONG64)buffer;
	info.targetAddress = (ULONG64)address;
	info.size = (ULONG64)size;
	if (DriverComm(CMD_READ_MEMORY, &info, sizeof(info))) {
		return TRUE;
	}
	return FALSE;
}

BOOLEAN DriverWriteProcessMemory(HANDLE pid, PVOID address, PVOID buffer, SIZE_T size)
{
	READWRITEINFO info;
	info.pid = (ULONG64)pid;
	info.buffer = (ULONG64)buffer;
	info.targetAddress = (ULONG64)address;
	info.size = (ULONG64)size;
	if (DriverComm(CMD_WRITE_MEMORY, &info, sizeof(info))) {
		return TRUE;
	}
	return FALSE;
}

BOOLEAN DriverQueryProcessMemory(HANDLE pid, PVOID address, PMEMORY_BASIC_INFORMATION memoryInfo)
{

	QUERYINFO info;
	info.pid = (ULONG64)pid;
	info.targetAddress = (ULONG64)address;
	info.memoryInfo = (ULONG64)memoryInfo;
	if (DriverComm(CMD_QUERY_MEMORY, &info, sizeof(info))) {
		return TRUE;
	}
	return FALSE;
}

BOOLEAN DriverProtectProcess(HANDLE pid)
{
	PROTECTEDPROCESSINFO info;
	info.pid = (ULONG64)pid;
	if (DriverComm(CMD_PROCESS_PROCESS, &info, sizeof(info))) {
		return TRUE;
	}
	return FALSE;
}

BOOLEAN DriverRemoteCall(HANDLE pid, PVOID shellCode, SIZE_T shellCodeSize)
{
	REMOTECALLINFO info;
	info.pid = (ULONG64)pid;
	info.shellCode = (ULONG64)shellCode;
	info.shellCodeSize = shellCodeSize;
	
	if (DriverComm(CMD_REMOTE_CALL, &info, sizeof(info))) {
		return TRUE;
	}
	return FALSE;
}

